package io.gatling.extensions.mysql.action

import com.typesafe.scalalogging.StrictLogging
import io.gatling.commons.stats.{KO, OK}
import io.gatling.commons.util.Clock
import io.gatling.core.CoreComponents
import io.gatling.core.action.{Action, BlockExit, ExitableAction}
import io.gatling.core.session.Session
import io.gatling.core.stats.StatsEngine
import io.gatling.core.util.NameGen
import io.gatling.extensions.mysql.builder._
import io.gatling.extensions.mysql.protocol.MysqlComponents

case class MysqlAction(
                        name: String,
                        operate: MysqlOperate,
                        mysqlComponents: MysqlComponents,
                        coreComponents: CoreComponents,
                        throttled: Boolean,
                        next: Action,
                      ) extends ExitableAction with NameGen with StrictLogging {

  import coreComponents._

  override def statsEngine: StatsEngine = coreComponents.statsEngine

  override def clock: Clock = coreComponents.clock

  override def execute(session: Session): Unit = {

    val start: Long = clock.nowMillis

    // 很奇怪，写法居然一样
    operate match {
      case Select(sql) => mysqlComponents.clients.get(session.userId).preparedQuery(sql(session).toOption.get).execute(asyncResult => if (asyncResult.succeeded()) success(session, start) else fail(session, asyncResult.toString, start))
      case Update(sql) => mysqlComponents.clients.get(session.userId).preparedQuery(sql(session).toOption.get).execute(asyncResult => if (asyncResult.succeeded()) success(session, start) else fail(session, asyncResult.toString, start))
      case Delete(sql) => mysqlComponents.clients.get(session.userId).preparedQuery(sql(session).toOption.get).execute(asyncResult => if (asyncResult.succeeded()) success(session, start) else fail(session, asyncResult.toString, start))
      case Insert(sql) => mysqlComponents.clients.get(session.userId).preparedQuery(sql(session).toOption.get).execute(asyncResult => if (asyncResult.succeeded()) success(session, start) else fail(session, asyncResult.toString, start))
    }
  }

  def fail(session: Session, message: String, start: Long): Unit = {
    statsEngine.logResponse(session.scenario, session.groups, name, start, clock.nowMillis, KO, Some(message), Some(message))
    throttler match {
      case Some(th) if throttled => th.throttle(session.scenario, () => next ! session)
      case None => next ! session
    }
  }

  def success(session: Session, start: Long): Unit = {
    statsEngine.logResponse(session.scenario, session.groups, name, start, clock.nowMillis, OK, None, None)
    throttler match {
      case Some(th) if throttled => th.throttle(session.scenario, () => next ! session)
      case None => next ! session
    }
  }
}
